//
//  transforms.cpp
//  yuzon
//
//  Created by Alvin on 2023/3/29.
//

#include "transforms.hpp"

using namespace cv;
using namespace std;

void transform_preds(const vector<float> &coords,
                     const Point2f &center,
                     const Size2f &scale,
                     const Size2f &output_size,
                     vector<float> &preds) {
  Mat trans = get_affine_transform(center, scale, 0.0f, output_size, 1);
  for (int p = 0; p < coords.size() / 2; ++p) {
    Mat trans_coord = affine_tranform(Point2f(coords[p * 2], coords[p * 2 + 1]), trans);
    preds[p * 3 + 1] = trans_coord.at<double>(0, 0);
    preds[p * 3 + 2] = trans_coord.at<double>(1, 0);
  }
}

Mat get_affine_transform(const Point2f &center,
                         const Size2f &scale,
                         float rot,
                         const Size2f &output_size,
                         int inv) {
  float src_w = scale.width;
  float dst_w = output_size.width;
  float dst_h = output_size.height;
  
  float rot_rad = M_PI * rot / 180;
  Point2f src_dir = get_dir(Point2f(0.0f, src_w * -0.5f), rot_rad);
  Point2f dst_dir = Point2f(0.0f, dst_w * -0.5f);
  
  Point2f src[3], dst[3];
  src[0] = center;
  src[1] = center + src_dir;
  dst[0] = Point2f(dst_w * 0.5f, dst_h * 0.5f);
  dst[1] = Point2f(dst_w * 0.5f, dst_h * 0.5f) + dst_dir;
  
  src[2] = get_3rd_point(src[0], src[1]);
  dst[2] = get_3rd_point(dst[0], dst[1]);
  
  Mat trans;
  
  if (inv == 1)
    trans = getAffineTransform(dst, src);
  else
    trans = getAffineTransform(src, dst);
  
  return trans;
}

Mat affine_tranform(const Point2f &pt, const Mat &t) {
  double data[] = {pt.x, pt.y, 1.0};
  Mat new_pt(3, 1, t.type(), data);
  new_pt = t * new_pt;
  return new_pt;
}

Point2f get_3rd_point(const Point2f &a, const Point2f &b) {
  Point2f direct = a - b;
  return b + Point2f(-direct.y, direct.x);
}

Point2f get_dir(const Point2f &src_point, float rot_rad) {
  float sn = sin(rot_rad);
  float cs = cos(rot_rad);
  
  Point2f src_result(0.0f, 0.0f);
  src_result.x = src_point.x * cs - src_point.y * sn;
  src_result.y = src_point.x * sn + src_point.y * cs;
  
  return src_result;
}

void get_final_preds(const vector<float> &heatmaps,
                     const vector<int64_t> &dims,
                     const vector<int64_t> &indexes,
                     const Point2f &center,
                     const Size2f scale,
                     vector<float> &preds,
                     bool dark) {
  int num_joints = int(dims[1]);
  int heatmap_h = int(dims[2]);
  int heatmap_w = int(dims[3]);
  int heatmap_s = heatmap_h * heatmap_w;
  
  vector<float> coords(num_joints * 2);
  
  for (int i = 0; i < num_joints; ++i) {
    // 接点对应热图在热图向量中的起始位置
    int heatmap_idx = i * heatmap_s;
    // 接点在对应热图中的位置
    int joint_idx = int(indexes[i]);
    
    // 接点在热图向量中的值
    float val = heatmaps[heatmap_idx + joint_idx];
    // 接点位置取模热图宽为px
    int px = joint_idx % heatmap_w;
    // 接点位置除以热图宽取整为py
    int py = floor(joint_idx / heatmap_w);
    
    preds[i * 3] = val;
    coords[i * 2] = px;
    coords[i * 2 + 1] = py;
    
    if (dark) {
      if (1 < px && px < heatmap_w - 2 && 1 < py && py < heatmap_h - 2) {
        vector<float> heatmap(heatmaps.begin() + heatmap_idx,
                              heatmaps.begin() + heatmap_idx + heatmap_s);
        dark_parse(heatmap, coords, heatmap_w, heatmap_h, px, py, i);
      }
    } else {
      // https://github.com/leoxiaobin/deep-high-resolution-net.pytorch/blob/master/lib/core/inference.py#L49
      if (0 < px && px < heatmap_w - 1) {
        float diff_x = (heatmaps[heatmap_idx + joint_idx + 1] -
                        heatmaps[heatmap_idx + joint_idx - 1]);
        coords[i * 2] += diff_x > 0.0f ? 0.25f : -0.25f;
      }
      
      if (0 < py && py < heatmap_h - 1) {
        float diff_y = (heatmaps[heatmap_idx + joint_idx + heatmap_w] -
                        heatmaps[heatmap_idx + joint_idx - heatmap_w]);
        coords[i * 2 + 1] += diff_y > 0.0f ? 0.25f : -0.25f;
      }
    }
  }
  
  transform_preds(coords, center, scale, Size2f(heatmap_w, heatmap_h), preds);
}

void dark_parse(const vector<float> &heatmap,
                vector<float> &coords,
                int w, int h, int px, int py, int i) {
  /*
   DARK postpocessing, Zhang et al. Distribution-Aware Coordinate
   Representation for Human Pose Estimation (CVPR 2020).
   1) offset = -(hassian.inv() * derivative)
   2) dx = (heatmap[x + 1] - heatmap[x - 1]) / 2
   3) dxx = (dx[x + 1] - dx[x - 1]) / 2
   4) derivative = Mat([dx, dy])
   5) hassian = Mat([[dxx, dxy], [dxy, dyy]])
   */
  Mat heatmap_mat = Mat(heatmap).reshape(0, h);
  heatmap_mat.convertTo(heatmap_mat, CV_32FC1);
  GaussianBlur(heatmap_mat, heatmap_mat, Size(3, 3), 0.0, 0.0);
  vector<float> hm = vector<float>(heatmap_mat.reshape(1, 1));
  
  float epsilon = 1e-10;
  
  // sample heatmap to get values in around target location
  float xy = log(fmax(hm[py * w + px], epsilon));
  float xr = log(fmax(hm[py * w + px + 1], epsilon));
  float xl = log(fmax(hm[py * w + px - 1], epsilon));
  float xr2 = log(fmax(hm[py * w + px + 2], epsilon));
  float xl2 = log(fmax(hm[py * w + px - 2], epsilon));
  float yd = log(fmax(hm[(py + 1) * w + px], epsilon));
  float yu = log(fmax(hm[(py - 1) * w + px], epsilon));
  float yd2 = log(fmax(hm[(py + 2) * w + px], epsilon));
  float yu2 = log(fmax(hm[(py - 2) * w + px], epsilon));
  float xryd = log(fmax(hm[(py + 1) * w + px + 1], epsilon));
  float xryu = log(fmax(hm[(py - 1) * w + px + 1], epsilon));
  float xlyd = log(fmax(hm[(py + 1) * w + px - 1], epsilon));
  float xlyu = log(fmax(hm[(py - 1) * w + px - 1], epsilon));
  
  // compute dx/dy and dxx/dyy with sampled values
  float dx = 0.5 * (xr - xl);
  float dy = 0.5 * (yd - yu);
  float dxx = 0.25 * ((xr2 - xy) - (xy - xl2));
  float dxy = 0.25 * ((xryd - xryu) - (xlyd - xlyu));
  float dyy = 0.25 * ((yd2 - xy) - (xy - yu2));
  
  // finally get offset by derivative and hassian, which combined by dx/dy and dxx/dyy
  if (dxx * dyy - dxy * dxy != 0) {
    float D[2] = {dx, dy};
    float M[2][2] = {dxx, dxy, dxy, dyy};
    
    Mat derivative(2, 1, CV_32F, D);
    Mat hassian(2, 2, CV_32F, M);
    Mat offset = -(hassian.inv() * derivative);
    
    coords[i * 2] += offset.at<float>(0, 0);
    coords[i * 2 + 1] += offset.at<float>(1, 0);
  }
}
